import requests
from exception import TraitError
url = "http://sre-api.gist.symantec.com/"
import os

class Trait:
    def __init__(self, md5):
        self.md5 = md5
        self.req = self._getTraitReq()
        self.traitmd5actionDict = {}
        self._getTraitMD5Action()
        
    # handles the request for traits for an md5
    # e.g. {u'size': 1001857, u'traits': [{u'trait2_object': u'Registry/Value', ... }
    def _getTraitReq(self):
        try: 
            req = requests.get(url + "sample/traits/" + self.md5 + ".json")
            return req.json()
        except ValueError:
            raise TraitError(self.md5 + (" produces bad Trait Request"))
    
    # stores as dict of {Trait:{Md5:details, Md5:details]}}
    # e.g {u'%TEMP%\\750829.exe': [{'10d4dbdfd14338617e0539ee84b22bf3': {u'trait2_object': .. }
    def _getTraitMD5Action(self):
        for item in self.req['traits']:
            detailPacker = {}
            detailPacker['trait_object'] = item['trait_object']
            detailPacker['trait'] = item['trait']
            detailPacker['trait_action'] = item['trait_action']
            md5action = {}
            md5action[self.md5] = detailPacker 
#             print md5action
            self.traitmd5actionDict[self._extractValue(item['trait'], item['trait_object'])] = md5action
            
            if item['trait2'] and item['trait'] is not item['trait2']:
                detailPacker = {}
                detailPacker['trait2_object'] = item['trait2_object']
                detailPacker['trait2'] = item['trait2']
                detailPacker['trait2_action'] = item['trait2_action']
                md5action = {}
                md5action[self.md5] = detailPacker
                self.traitmd5actionDict[self._extractValue(item['trait2'], item['trait2_object'])] = md5action
            
    # value to compare is extracted from the trait string depending on the type of trait
    def _extractValue(self, traitStr, traitType):
        # early return for String, Detection and Static trait type
        if 'String' in traitType  or 'Detection' in traitType or 'Static' in traitType:
            value = traitStr
        
        # used to extract the parameters in API and Window trait type
        # params will be connected together as "value"
        elif 'API' in traitType or 'Window' in traitType:
            if ',' in traitStr:
                value = ""
                temp = traitStr.split(',')
                for param in temp:
                    if '=' in param:
                        value += param.split('=')[0]
                    else:
                        value = traitStr
            else:
                value = traitStr
                  
        # used for extracting value/key/filename from a registry and file trait object
        # e.g. %TEMP%\TMP33A~1.BAT -> TMP33A~1
        elif 'Registry' in traitType:
            value = traitStr
            if '\\' in value:
                temp = value.split('\\')
                value = temp.pop()
                if '\"\"\"' in value:
                    value = temp.pop()
                value = os.path.splitext(value)[0]
            
        # used for extracting filename from file trait type
        elif 'File' in traitType or 'Directory' in traitType:
            value = traitStr
            if '%\\' in value:   # removes directory path
                value = value.split('%\\').pop()
            if '\\' in value:
                value = value.split('\\').pop()  # removes file extension
            value = os.path.splitext(value)[0]

        # used for extracting filename from Network/http trait type
        elif 'HTTP' in traitType:
            value = traitStr
            if '/' in value:
                value = value.split('/').pop()
            value = os.path.splitext(value)[0]
        
        # used for extracting network name from Network trait type
        elif traitType == 'Network':
            value = traitStr.lstrip('www.')
            temp, ext = os.path.splitext(value) # removes file extension
            ext = ext.lstrip('.')
            if len(ext) <= 3 and ext.isalpha(): # check if really extension
                value = temp

        # used for extracting mutex from Mutex trait type
        elif 'Mutex' in traitType:
            value = traitStr
            if '\\' in value:
                value = value.split('\\').pop()  # get the last value only
       
        # used for extracting version from Version trait type
        elif 'Version' in traitType:
            value = traitStr
            temp, ext = os.path.splitext(value) # removes file extension
            ext = ext.lstrip('.')
            if len(ext) <= 3 and ext.isalpha():
                value = temp
                
        # take param as it is
        else:
            value = traitStr
            
        return value
            